#include <Utility/CLib.h>
#include <Memory/Heap.h>
#include <Memory/Paging.h>
#include <IntExcept/Debug.h>
#include <Utility/intrinsic.h>
/*
@brief 在堆中分配内存空间
@param `lpHeap` 指向堆描述符表的指针
@param `dwFlags` 选项
@param `dwBytes` 需要分配的字节数
@return 堆内存块地址
*/
LPVOID HeapAllocate(LPArenaDescriptorTable lpHeap, DWORD dwFlags, DWORD dwBytes)
{
    BYTE bPower;

    if (dwBytes == 0) dwBytes = 16;
    __asm
    {
        mov eax, 1 // 最终结果由 eax 移位得到
        mov edx, dwBytes
        bsr ecx, edx // ecx 保存 most significant set bit 出现的位置
        btr edx, ecx // most significant set bit 清空
        test edx, 0xFFFFFFFF // 检测 edx 中是否还有非 0 的位
        jz KernelHeapAllocate_Round_Finish // 只剩下 0，说明是 2 的整数次幂
        inc ecx // 还有非 0 的位，不是 2 的整数次幂，移位次数需要加 1
    KernelHeapAllocate_Round_Finish:
        shl eax, cl
        mov bPower, cl // 2 的多少次方
        mov dwBytes, eax // 更新 dwBytes 为不小于其的最小 2 的整数次幂
    }
    /* dwBytes 扩大为 2 的整数次幂 */
    int i = 0;
    LPVOID lpRet = NULL;
    if (dwBytes < 256) /* 从 Small Arena 分配 */
    {
        while (TRUE)
        {
            if (lpHeap[i].wGranularity == dwBytes && lpHeap[i].wAvailableUnits != 0) /* 类型匹配、有剩余可用 */
            {
                /* 校验：wcAvailable 必须和位图中的可用位数目一致，否则必定出错 */
                if (lpHeap[i].wGranularity == 16)
                {
                    KERNEL_ASSERT((*(LPDWORD)(lpHeap[i].lpArena) & 0x3) == 0, "First two bits in 16 B arena's bitmap not zero.\r\n"); /* 16 B Arena 位图首 2 位必为 0 */
                }
                else 
                {
                    KERNEL_ASSERT((*(LPDWORD)(lpHeap[i].lpArena) & 0x1) == 0, "First bit in arena's bitmap not zero.\r\n");
                } /* 其余种类的小型 Arena 位图第一位必为 0 */
                WORD wAvailable = 0;
                for (int j = 0; j < ARENA_BITMAP_SIZE_IN_DWORD(dwBytes); j++)
                {
                    wAvailable += __popcnt(((LPDWORD)(lpHeap[i].lpArena))[j]);
                }
                KERNEL_ASSERT(wAvailable == lpHeap[i].wAvailableUnits, "Inconsistent number of bits available in arena bitmap with wAvailableUnits field in arena descriptor.\r\n");
                DWORD dwOffset = AquireAvailableBit(ARENA_BITMAP_SIZE_IN_DWORD(dwBytes) ,lpHeap[i].lpArena);
                KERNEL_ASSERT(dwOffset != MEMORY_BITMAP_INVALID_OFFSET, "Physical memory pool runs off.\r\n"); /* wAvailableUnits 必定与位图中可用位的数目一致，否则出错 */
                lpHeap[i].wAvailableUnits--;
                lpRet = (LPBYTE)(lpHeap[i].lpArena) + (dwOffset << bPower); /* 此处使用 bPower 提速，防止使用乘法器 */
                break; /* 分配完毕 */
            }
            else if (lpHeap[i].wGranularity == ARENA_NONE) /* 空 Arena 描述符 */
            {
                lpHeap[i].wGranularity = dwBytes;
                InitializeArena(&lpHeap[i]); /* 初始化该描述符，分配页面 */
                DWORD dwOffset = AquireAvailableBit(ARENA_BITMAP_SIZE_IN_DWORD(dwBytes) ,lpHeap[i].lpArena);
                KERNEL_ASSERT(dwOffset != MEMORY_BITMAP_INVALID_OFFSET, "Physical memory pool runs off.\r\n"); /* wAvailableUnits 必定与位图中可用位的数目一致，否则出错 */
                lpHeap[i].wAvailableUnits--;
                lpRet = (LPBYTE)(lpHeap[i].lpArena) + (dwOffset << bPower); /* 此处使用 bPower 提速，防止使用乘法器 */
                break; /* 分配完毕 */
            }
            else if (lpHeap[i].wGranularity == ARENA_MORE) /* 当前 Arena 描述符指向另一个 Arena */
            {
                lpHeap = lpHeap[i].lpArena; /* 继续处理下个 Arena 描述符 */
                i = 0;
                continue;
            }
            i++;
        }
    }
    else if (dwBytes < 4096) /* 可以通过基本的 Arena 分配 */
    {
        while (TRUE)
        {
            
            if (lpHeap[i].wGranularity == dwBytes && lpHeap[i].wAvailableUnits != 0) /* 类型匹配、有剩余可用 */
            {
                DWORD dwBitmap = lpHeap[i].wAvailableUnits; /* AquireAvailableBit 操作的单位为 DWORD */
                DWORD dwOffset = AquireAvailableBit(1, &dwBitmap); /* 查询位图中的可用位 */
                KERNEL_ASSERT(dwOffset != MEMORY_BITMAP_INVALID_OFFSET, "Physical memory pool runs off.\r\n"); /* wAvailableUnits 必定与位图中可用位的数目一致，否则出错 */
                lpHeap[i].wAvailableUnits = (WORD)dwBitmap; /* 更新修改后的位图 */
                lpRet = (LPBYTE)(lpHeap[i].lpArena) + (dwOffset << bPower); /* 此处使用 bPower 提速，防止使用乘法器 */
                break; /* 分配完毕 */
            }
            else if (lpHeap[i].wGranularity == ARENA_NONE) /* 空 Arena 描述符 */
            {
                lpHeap[i].wGranularity = dwBytes;
                InitializeArena(&lpHeap[i]); /* 初始化 Arena */
                DWORD dwBitmap = lpHeap[i].wAvailableUnits; /* 位图大小扩展为 1 个 DWORD */
                DWORD dwOffset = AquireAvailableBit(1 ,&dwBitmap); /* 查询位图中的可用位 */
                KERNEL_ASSERT(dwOffset != MEMORY_BITMAP_INVALID_OFFSET, "Physical memory pool runs off.\r\n"); /* wAvailableUnits 必定与位图中可用位的数目一致，否则出错 */
                lpHeap[i].wAvailableUnits = (WORD)dwBitmap; /* 更新修改后的位图 */
                lpRet = (LPBYTE)(lpHeap[i].lpArena) + (dwOffset << bPower); /* 此处使用 bPower 提速，防止使用乘法器 */
                break; /* 分配完毕 */
            }
            else if (lpHeap[i].wGranularity == ARENA_MORE) /* 当前 Arena 描述符指向另一个 Arena */
            {
                lpHeap = lpHeap[i].lpArena; /* 继续处理下个 Arena 描述符 */
                i = 0;
                continue;
            }
            i++;
        }
    }
    else /* 需要连续多个 4096 B Arena，超大型堆 */
    {
        while (TRUE) /* 超大型堆仅在需要时创建 */
        {
            if (lpHeap[i].wGranularity == ARENA_NONE) /* 空 Arena 描述符 */
            {
                lpHeap[i].wGranularity = 4096;
                lpHeap[i].wAvailableUnits = dwBytes / 4096;
                InitializeArena(&lpHeap[i]); /* 初始化 Arena */
                lpRet = lpHeap[i].lpArena;
                break; /* 分配完毕 */
            }
            else if (lpHeap[i].wGranularity == ARENA_MORE) /* 当前 Arena 描述符指向另一个 Arena */
            {
                lpHeap = lpHeap[i].lpArena; /* 继续处理下个 Arena 描述符 */
                i = 0;
                continue;
            }
            i++;
        }
    }
    /* 根据 dwFlags 对分配得到的堆进行相应处理 */
    if (dwFlags == HEAP_ZERO_MEMORY)
    {
        memset(lpRet, 0, dwBytes);
    }
    return lpRet;
}

